'''
对接本地化通义千问大模型，实现自主知识库内容的增强查询
'''
# from langchain.embeddings.openai import OpenAIEmbeddings
# from langchain.vectorstores import Chroma
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
# from langchain import OpenAI
# 这种方式不推荐
# from langchain.document_loaders import DirectoryLoader
# 根据文档查找新的导入方式
from langchain_community.document_loaders import DirectoryLoader
from langchain.chains import RetrievalQA
from ChatQwen import ChatQwen
from langchain_community.embeddings import HuggingFaceEmbeddings
import os
# from dotenv import load_dotenv

# load_dotenv("config.env")
# embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
embeddings_model_name = "shibing624/text2vec-base-chinese"

# 加载文件夹中的所有txt类型的文件
loader = DirectoryLoader('./docs', glob='*.txt')
# 将数据转成 document 对象，每个文件会作为一个 document
documents = loader.load()

# 初始化加载器
text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
# 切割加载的 document
split_docs = text_splitter.split_documents(documents)

# 首先加入以下代码，修改下载镜像
os.environ['HF_ENDPOINT']='https://hf-mirror.com'
# 初始化 openai 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
# 将 document 通过 openai 的 embeddings 对象计算 embedding 向量信息并临时存入 Chroma 向量数据库，用于后续匹配查询
docsearch = Chroma.from_documents(split_docs, embeddings)

# 创建问答对象
qa = RetrievalQA.from_chain_type(llm=ChatQwen(), chain_type="stuff", retriever=docsearch.as_retriever(), return_source_documents=True)
# 进行问答
result = qa({"query": "怎么评估自主开发工作包？"})
print(result)